import argparse
import json
import logging
import time

import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.requests import Request

from lagent.schema import AgentMessage
from lagent.utils import load_class_from_string


class AgentAPIServer:
    """代理 API 服务器，提供 HTTP 接口"""

    def __init__(self,
                 config: dict,
                 host: str = '127.0.0.1',
                 port: int = 8090):
        """
        初始化 API 服务器。

        Args:
            config: 代理配置
            host: 服务器主机地址
            port: 服务器端口
        """
        # 创建 FastAPI 应用
        self.app = FastAPI(docs_url='/')
        
        # 配置 CORS 中间件，允许跨域请求
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=['*'],  # 允许所有来源
            allow_credentials=True,  # 允许携带凭证
            allow_methods=['*'],  # 允许所有方法
            allow_headers=['*'],  # 允许所有头部
        )

        # 动态加载代理类并实例化
        cls_name = config.pop('type')
        python_path = config.pop('python_path', None)
        cls_name = load_class_from_string(cls_name, python_path) if isinstance(
            cls_name, str) else cls_name
        self.agent = cls_name(**config)

        # 设置路由并启动服务器
        self.setup_routes()
        self.run(host, port)

    def setup_routes(self):
        """设置 API 路由"""

        def heartbeat():
            """健康检查接口"""
            return {'status': 'success', 'timestamp': time.time()}

        async def process_message(request: Request):
            """
            处理聊天消息接口。
            
            Args:
                request: HTTP 请求对象
                
            Returns:
                代理响应消息
                
            Raises:
                HTTPException: 处理失败时抛出
            """
            try:
                body = await request.json()
                # 解析消息
                message = [
                    m if isinstance(m, str) else AgentMessage.model_validate(m)
                    for m in body.pop('message')
                ]
                # 调用代理处理消息
                result = await self.agent(*message, **body)
                return result
            except Exception as e:
                logging.error(f'Error processing message: {str(e)}')
                raise HTTPException(
                    status_code=500, detail='Internal Server Error')

        def get_memory(session_id: int = 0):
            """
            获取会话记忆接口。
            
            Args:
                session_id: 会话ID
                
            Returns:
                会话状态数据
                
            Raises:
                HTTPException: 处理失败时抛出
            """
            try:
                result = self.agent.state_dict(session_id)
                return result
            except KeyError:
                raise HTTPException(
                    status_code=404, detail="Session ID not found")
            except Exception as e:
                logging.error(f'Error processing message: {str(e)}')
                raise HTTPException(
                    status_code=500, detail='Internal Server Error')

        # 注册路由
        self.app.add_api_route('/health_check', heartbeat, methods=['GET'])
        self.app.add_api_route(
            '/chat_completion', process_message, methods=['POST'])
        self.app.add_api_route(
            '/memory/{session_id}', get_memory, methods=['GET'])

    def run(self, host='127.0.0.1', port=8090):
        """启动服务器"""
        logging.info(f'Starting server at {host}:{port}')
        uvicorn.run(self.app, host=host, port=port)


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Async Agent API Server')
    parser.add_argument('--host', type=str, default='127.0.0.1')
    parser.add_argument('--port', type=int, default=8090)
    parser.add_argument(
        '--config',
        type=json.loads,
        required=True,
        help='JSON configuration for the agent')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    # 配置日志并启动服务器
    logging.basicConfig(level=logging.INFO)
    args = parse_args()
    AgentAPIServer(args.config, host=args.host, port=args.port)
